// K-Shortest_Paths_Optimization-2.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"

#include <iostream>
#include <boost/property_map/vector_property_map.hpp>
#include <boost/graph/dag_shortest_paths.hpp>
#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <boost/graph/adjacency_list.hpp>
#include "boost/filesystem.hpp"
#include "boost/lexical_cast.hpp"
#include <boost/regex.hpp>
#include <boost/iostreams/device/file.hpp> 
#include <boost/iostreams/stream.hpp> 
#include <boost/graph/copy.hpp>
#include <vbl/img/ImageDraw.h>
#include <vbl/img/ImageCopy.h>
#include <vbl/img/ImageUtils.h>
#include <vbl/img/io/ImageIO.h>
#include <vbl/img/proc/Demosaic.h>
#include "vbl/util/window.h"
#include <vbl/img/ImageTransform.h>
#include <boost/progress.hpp>
#include <vbl/util/Stopwatch.h>
#include "HelperFunctions.h"

namespace bf = boost::filesystem;
using namespace std;
using namespace boost;
using namespace vbl;

typedef adjacency_list<vecS, vecS, directedS, 
		property<vertex_distance_t, double>, property<edge_weight_t, double> > graph_t;
typedef graph_traits<graph_t>::vertex_descriptor vertex_t;
typedef graph_traits<graph_t>::edge_descriptor edge_t;
typedef property_map<graph_t, vertex_distance_t>::type distance_map_t;
typedef property_map<graph_t, edge_weight_t>::type weight_map_t;
typedef property_map<graph_t, vertex_index_t>::type index_map_t;
typedef graph_traits<graph_t>::vertex_iterator vertex_iter;
typedef graph_traits<graph_t>::edge_iterator edge_iter;

bool VERBOSE;
int VSOURCE;

struct location {
	int ts;
	int x;
	int y;
};

struct rectangle {
	bool isVisible;
	int xmin;
	int xmax;
	int ymin;
	int ymax;
};

struct cutOutInfo {
	rectangle r;
	int pathNr;
	img::ImageRGB imageCutOut;
};

void getNeighbours(int location, int numOfLocations, int vsink , vector<int> & neighbours , int pomX ){
	// location is not yet correct, location + 1 is the final location in the graph!
	// TODO: change this!
	// the location in a timestep
	int location_in_timestep = location % numOfLocations;

	bool top = true;
	bool bottom = true;
	bool left = true;
	bool right = true;

	// check if the location is on an edge of the room
	if( location_in_timestep % pomX == 0 ) { 
		// node is on left side of grid
		left = false;
	}
	if( location_in_timestep % pomX == (pomX - 1) ) {
		// node is on right side of grid
		right = false;
	}
	if( location_in_timestep < pomX ) {
		// node is on top side of grid
		top = false;
	}
	if( location_in_timestep >= ( numOfLocations - pomX ) ) {
		// node is on bottom side of grid
		bottom = false;
	}

	if( location > vsink - numOfLocations - 2 ) {
		if( !left || !right || !top || !bottom ) {
			neighbours.push_back(VSOURCE);
		}
		// TODO: vsink should be total_locations
		// location is situated the last timestep,
		// so it only needs to be connected to vsink
		// "-2" because of the extra locations vsource and vsink
		neighbours.push_back(vsink);
		return;
	}

	// correct location, vsource = 0 so all locations move up by 1
	location = location + 1;
	neighbours.push_back(location + numOfLocations);											// same location
	if( left ) {
		neighbours.push_back(location - 1 + numOfLocations);						// left
	}
	if( right ) {
		neighbours.push_back(location + 1 + numOfLocations);						// right
	}
	if( top ) {
		neighbours.push_back(location - pomX + numOfLocations);			// top
		if( left ) {
			neighbours.push_back(location - pomX - 1 + numOfLocations);	// top left
		}
		if( right ) {
			neighbours.push_back(location - pomX + 1 + numOfLocations);	// top right
		}
	}
	if( bottom ) {
		neighbours.push_back(location + pomX + numOfLocations);			// bottom
		if( left ) {
			neighbours.push_back(location + pomX - 1 + numOfLocations);	// bottom left
		}
		if( right ) {
			neighbours.push_back(location + pomX + 1 + numOfLocations);	// bottom right
		}
	}
	if( !left || !right || !top || !bottom ) {
			neighbours.push_back( VSOURCE );
			neighbours.push_back( vsink );
	}
}

void readFile( int i , vector<double> & vertices , vector<double> & proba , string probaPath ) {
	// find and read file
	string frameNumber = lexical_cast<string>(i);
	//string fileName = probaPath + "/" + frameNumber + "-proba.dat";	// MM
	string fileName = probaPath + "/proba-f" + frameNumber + ".dat";	// CAS
	ifstream file(fileName.c_str());

	if(!file.is_open()){
		cout << "Error opening proba file:" << endl;
		cout << fileName << endl;
		exit(1);
	} else {
		string line;
		while(getline(file, line)){
			// find 'space'-character, value behind it is probability, value before it is location
			for( unsigned int j = 0 ; j < line.size() ; j++ ){
				char c = line.at(j);
				if(c == ' '){
					line = line.substr(j+1,line.size());
					break;
				}
			}
			double probability = lexical_cast<double>(line);

			proba.push_back(probability);

			// TEMP
			// TODO: REMOVE THIS. ONLY FOR GT DATA
			if( probability == 1 ) {
				probability = 0.9999;
			}
			if( probability == 0 ) {
				probability = 0.0001;
			}
			// TEMP
			double cost = -log(probability / (1-probability));

			vertices.push_back(cost);
		}
		file.close();
	}
}

double path_cost( vector<int> path , graph_t & g , weight_map_t weight_map ) {
	double w = 0.0;
	for( unsigned int j = 0 ; j < path.size() ; j++ ) {
		if( j != 0 ) {
			edge_t e;
			bool found;
			tie(e, found) = edge(path[j-1],path[j], g);
			w += get(weight_map, e);
		}
	}
	return w;
}


double paths_cost(vector<vector<int>> paths, graph_t & g ) {
	weight_map_t weight_map = get( edge_weight, g );
	double w = 0.0;
	for( unsigned int i = 0 ; i < paths.size() ; i++ ){
		vector<int> path = paths[i];
		w += path_cost( path , g , weight_map );
	}
	return w;
}

void extend_graph_path( graph_t & g , vector<int> path , int vsink , map<int,int> & vertexPairs ) {
	weight_map_t weight_map = get( edge_weight , g );
	vector<int> new_path;

	for ( unsigned int i = 0 ; i < path.size() ; i++ ) {
		new_path.push_back( path[i] );
		if ( path[i] != VSOURCE && path[i] != vsink ) {
			// add new vertex
			vertex_t v = add_vertex( g );
			new_path.push_back( v );
			vertexPairs.insert(make_pair(v,path[i]));
			// connect all out edges of path[i] to new vertex
			graph_traits<graph_t>::out_edge_iterator out_it, out_it_end;
			tie( out_it, out_it_end ) = out_edges( path[i] , g );
			vector<vertex_t> toRemove;
			double cost = get( weight_map , *out_it );
			for ( ; out_it != out_it_end ; out_it++ ) {
				vertex_t t = target( *out_it , g );
				add_edge( v , t , cost, g );
				toRemove.push_back( t );
			}
			
			// remove all out edges from path[i]
			for ( unsigned int j = 0 ; j < toRemove.size() ; j++ ) {
				remove_edge( path[i] , toRemove[j] , g );
			}
			
			// connect path[i] to new vertex with cost 0
			add_edge( path[i] , v , 0 , g );
		}
	}

	// for all vertices in the path, reverse direction and negate cost
	// TODO: could put size -1! this does not turn around the vertex to the sink node, check if this is what we want
	for ( unsigned int i = 1 ; i < new_path.size() ; i++ ) {
		edge_t e;
		bool found;
		tie( e , found ) = edge( new_path[i-1] , new_path[i], g );
		if( found ) {
			double cost = get( weight_map , e );
			remove_edge( new_path[i-1] , new_path[i], g );
			add_edge( new_path[i] , new_path[i-1] , -cost , g );
		} else {
			cout << "Edge not found while reversing cost and direction of a path" << endl;
			cout << "Edge going from " << new_path[i-1] << " to " << new_path[i] << endl;
			exit(1);
		}
	}
}

void extend_graph( graph_t & g , vector<vector<int>> paths , int vsink, map<int,int> & vertexPairs) {
	if( VERBOSE ) {
		cout << "...(Re-)Extending Graph..." << endl;
	}
	for ( unsigned int i = 0 ; i < paths.size() ; i++ ) {
		extend_graph_path( g , paths[i] , vsink , vertexPairs );
	}
	if( VERBOSE ) {
		cout << "...Done..." << endl;
	}
}

void transform_edge_cost( graph_t & g , map<pair<int,int>, double> weights , vector_property_map<double> & distance_map , unsigned int vsink , unsigned int totalLocations , map<int,int> vertexPairs ) {
	if ( VERBOSE ) {
		cout << "...Transforming Edges..." << endl;
	}
	//weight_map_t weight_map = get( edge_weight , g_prev );
	edge_iter ei, ei_end;
	tie( ei , ei_end ) = edges( g );
	for( ; ei != ei_end ; ei++ ) {
		vertex_t s = source( *ei , g );
		vertex_t t = target( *ei , g );

		pair<int,int> key(s,t);
		map<pair<int,int>, double>::iterator wi = weights.find(key);
		double cost;
		if( wi != weights.end() ) {
			cost = (*wi).second;
		} else {
			if( t == 0 ) {
				pair<int,int> key(t,s);
				cost = -weights[key];
			} else if( t >= totalLocations ) {
				vertex_t pairedVertex = vertexPairs[t];
				pair<int,int> key(pairedVertex,s);
				cost = -weights[key];
			} else if( s >= totalLocations ) {
				vertex_t pairedVertex = vertexPairs[s];
				if ( pairedVertex == t ) {
					cost = 0;
				} else {
					pair<int,int> key(pairedVertex,t);
					cost = weights[key];
				}
			} else {
				cost = get( edge_weight , g , *ei );
			}
		}
//		double cost = weights[key];//get( weight_map , *ei ); //get( edge_weight , g_prev , *ei );
		double distanceSource;
		double distanceTarget;

		distanceSource = distance_map[s];
 
		if ( t >= totalLocations ) {
			distanceTarget = distanceSource + cost;
			distance_map[t] = distanceTarget;
		} else {
			distanceTarget = distance_map[t];
		}
		
		if ( distanceTarget > distanceSource + cost ) {
			distanceTarget = distanceSource + cost;
			distance_map[t] = distanceTarget;
		}
		
		/*if ( distanceSource == distanceTarget ) {
			cost = 0;
		}*/
		double new_cost = cost + distanceSource - distanceTarget;
		if( new_cost < 0 && new_cost > -0.001 ) {
			new_cost = 0;
		}
		put( edge_weight , g , *ei , new_cost );
	}

	if ( VERBOSE ) {
		cout << "...Done..." << endl;
	}
}

vector<int> find_and_remove_competing_path(vector<vector<int>> & paths, int vertex , unsigned int & continueLocation ){
	vector<int> toReturn;
	int path_location;
	for( unsigned int i = 0 ; i < paths.size() ; i++ ) {
		vector<int> path = paths[i];
		for( unsigned int j = 0 ; j < path.size() ; j++ ) {
			int toCompare = path[j];
			if( vertex == toCompare ) {
				toReturn = paths[i];
				path_location = i;
				continueLocation = j;
				break;
			}
		}
	}
	paths.erase(paths.begin() + path_location);
	return toReturn;
}

int locationOf( int toSearch , vector<int> path ) {
	for( unsigned int i = 0 ; i < path.size() ; i++ ) {
		if( path[i] == toSearch ) {
			return i;
		}
	}
	return -1;
}

void find_leave_point( vector<int> path , vector<int> competing_path , int hit_location , unsigned int & leave_location , unsigned int & split_location ) {
	split_location = locationOf( path[hit_location] , competing_path );
	while( true ) {
		split_location -= 1;
		leave_location = locationOf( competing_path[split_location] , path );
		if ( leave_location == -1 ) {
			leave_location = locationOf( competing_path[split_location+1] , path ) + 2;
			break;
		}
	}
}

/*void find_exit_points( vector<int> path , vector<int> competingPath , int entryPoint , int & exitPointpath , int & exitPointCompetingPath ) {
	
}

void interlace_and_augment( vector<int> path , vector<vector<int>> & paths , map<int,int> & vertexPairs , int totalLocations ) {
	for( unsigned int i = 0 ; i < path.size() ; i++ ) {
		if( path[i] > totalLocations ) {
			// not in orignal graph, so must be one of the edges going backwards through an older path
			vector<int> competingPath = find_and_remove_competing_path(paths, path[hit_location]);
			int entryPointPath = i - 1;
			int exitPointPath;
			int exitPointCompetingPath;
			find_exit_points( path , competingPath , entryPointPath , exitPointPath , exitPointCompetingPath );
		}
	}
}*/


void interlace_and_augment( vector<int> path , vector<vector<int>> & paths, graph_t & g_backup) {
	// apply to path
	unsigned int hit_location = 0;
	unsigned int leave_location = 0;
	// applies to crossing path
	unsigned int split_location = 0;
	unsigned int continueLocation = 0;
	vector<int> competing_path;
	bool interlace_found = false;
	for( unsigned int i = 0 ; i < path.size() ; i++ ) {
		if( i != 0 ) {
			edge_t e;
			bool found;
			tie(e, found) = edge(path[i-1],path[i], g_backup);
			if( !found ) {
				// not in orignal graph, so must be one of the edges going backwards
				hit_location = i-1;
				competing_path = find_and_remove_competing_path(paths, path[hit_location] , continueLocation );
				interlace_found	= true;
				find_leave_point( path , competing_path , hit_location , leave_location , split_location );
				break;
			}
		}
	}
	if( interlace_found ) {
		vector<int> new_path_1;
		vector<int> new_path_2;
		for( unsigned int i = 0 ; i < hit_location ; i++ ) {
			new_path_1.push_back(path[i]);
		}
		for( unsigned int i = continueLocation ; i < competing_path.size() ; i++ ){
			new_path_1.push_back(competing_path[i]);
		}
		for( unsigned int i = 0 ; i <= split_location ; i++ ) {
			new_path_2.push_back(competing_path[i]);
		}
		for( unsigned int i = leave_location ; i < path.size() ; i++ ){
			new_path_2.push_back(path[i]);
		}
		//paths.push_back( new_path_1 );
		paths.push_back( new_path_1 );
		interlace_and_augment( new_path_2 , paths , g_backup );
	} else {
		paths.push_back( path );
	}
}

void printPath( vector<int> path , int numOfLocations , int pomX ) {
	for( unsigned int j = 0 ; j < path.size() ; j++ ) {
		int loc;
		int locx;
		int locy;  
		int ts;

		if ( path[j] == VSOURCE ) {
			locx = 0;
			locy = 0;
			ts = -1;
		} else {
			loc = ( ( path[j] - 1 ) % numOfLocations ) + 1;
			locx = ((loc-1) % pomX )+1;
			locy = (( loc - locx ) / pomX )+1;
			ts = (int) floor( ((double) path[j]-1) / numOfLocations );
		}
		cout << "(" << ts << " : " << locx << ", " << locy << ")" << ", ";
	}
	cout << endl;
}

void printPaths( vector<vector<int>> paths , int numOfLocations , int pomX ) {
	for( unsigned int i = 0 ; i < paths.size() ; i++ ) {
		printPath( paths[i] , numOfLocations , pomX );
	}
	cout << paths.size() << " paths found." << endl;
}

void reconstructPath( vector<size_t> predecessorMap , int source , int vsink , vector<int> & path ) {
	int v = vsink;
	while( v != source ) {
		path.insert( path.begin(), v );
		int v_temp = predecessorMap[v];
		if( v == v_temp ) {
			cout << "Node is its own predecessor" << endl;
			exit(0);
		} else {
			v = v_temp;
		}
	}
	path.insert( path.begin(), source );
	if ( source != VSOURCE ) {
		path.insert( path.begin(), VSOURCE );
	}
}

vector<int> find_shortest_path_dijkstra( graph_t & g , int source , int vsink , distance_map_t & distance_map ) {
	if( VERBOSE ) {
		cout << "Performing Dijkstra's shortest paths algorithm" << endl;
	}
	vector<size_t> predecessorMap(num_vertices(g));
       
	distance_map = get( vertex_distance , g );
	weight_map_t weight_map = get( edge_weight , g );
	index_map_t index_map = get( vertex_index , g );
	
	// initialize variables for dag_shortest_paths
	default_dijkstra_visitor vis;
	closed_plus<double> combine;
	less<double> compare;
    
	dijkstra_shortest_paths(g, source, &predecessorMap[0], distance_map, weight_map, index_map, compare, combine, (numeric_limits<double>::max)(), 0, vis);
	
	// find shortest path by tracing back from vsink
	vector<int> path;
	reconstructPath( predecessorMap , source , vsink , path );
	
	return path;
}

vector<int> find_shortest_path_dag( graph_t & g , int source , int vsink , distance_map_t & distance_map ) {
	// find graph properties
	distance_map = get( vertex_distance , g );
	weight_map_t weight_map = get( edge_weight, g );
	vector<default_color_type> color(num_vertices(g));
	vector<size_t> predecessor_map(num_vertices(g));
	
	// initialize variables for dag_shortest_paths
	default_dijkstra_visitor vis;
	closed_plus<double> combine;
	less<double> compare;
	
	// perform dag_shortest_paths
	dag_shortest_paths( g , source , distance_map, weight_map, &color[0], &predecessor_map[0], vis, compare, combine, (numeric_limits<int>::max)(), 0);

	// find shortest path by tracing back from vsink
	vector<int> path;
	reconstructPath( predecessor_map , source , vsink , path );
	
	return path;
}

void show_gt( img::ImageRGB & im , int frame , string path , int scale) {
	string frameNumber = lexical_cast<string>(frame);
	string fileName = path + "/" + frameNumber + "-gt.dat";
	ifstream file(fileName.c_str());

	if(!file.is_open()){
		cout << "Error opening gt file." << endl;
		return;
	} else {
		string line, x, y;
		int k = 0;
		while(getline(file, line)){
			// find 'space'-character, value behind it is y, value before it is x
			for( unsigned int j = 0 ; j < line.size() ; j++ ){
				char c = line.at(j);
				if(c == ' '){
					x = line.substr(0, j);
					y = line.substr(j+1,line.size());
					break;
				}
			}
			// rounding values
			int x_int = int( (lexical_cast<double>(x) * scale) + 0.5 );
			int y_int = int( (lexical_cast<double>(y) * scale) + 0.5 );
			img::drawCross ( im , y_int+1 , x_int+1 , img::Triple<>(255,0,0) , 5 , 1 );
		}
		file.close();
	}
}

vector<double> readFiles( int beginFrame , int numOfFrames , string path , vector<vector<double>> & probabilities ) {
	vector<double> vertices;
	int endFrame = beginFrame + numOfFrames;
	int percentage = 10; // percentage used for printing a dot for reading progress

	if( VERBOSE ) {
		cout << "...Reading files..." << endl;
	}
	
	for( int i = beginFrame ; i < endFrame ; i++){
		if( VERBOSE ) {
			if ( ( ( ( i - beginFrame + 1 ) / (double) numOfFrames ) * 100.0 ) >= percentage ) {
				cout << ".";
				percentage += 10;
			}
		}
		vector<double> proba;
		readFile(i, vertices, proba, path);
		probabilities.push_back(proba);
	}

	if( VERBOSE ) {
		cout << endl;
		cout << "...Done..." << endl;
	}

	return vertices;
}

graph_t createGraph( vector<double> vertices , int vsink , unsigned int numOfLocations , graph_t & g , int pomX , bool firstGraph , vector<vector<location>> paths , int batchSize , int lastTimeStep) {	
	if ( VERBOSE ) {
		cout << "...Creating Graph..." << endl;
	}
	
	if ( firstGraph ) {
		// connect vsource to all nodes in first timestep
		for( unsigned int i = 1 ; i <= numOfLocations ; i++ ){
			add_edge( VSOURCE , i , 0 , g );
		}
	} else {
		for( unsigned int i = 0 ; i < paths.size() ; i++ ) {
			vector<location> path = paths[i];
			location loc = path[path.size()-2];
			if ( loc.ts == lastTimeStep ) {
				int l = (loc.y-1)*pomX + loc.x;
				add_edge( VSOURCE , l , 0 , g );
			}
		}
	}
	for( unsigned int i = 0 ; i < vertices.size() ; i++ ){
		double cost = vertices[i];
		vector<int> neighbours;
		getNeighbours( i , numOfLocations , vsink , neighbours , pomX );
		for( unsigned int j = 0 ; j < neighbours.size() ; j++ ) {
			// TODO: ugly coding, +1 is to correct for nodes beginning at 0 (vsource == 0)
			if( neighbours[j] == VSOURCE ) {
				if ( i+1 > numOfLocations ) {
					add_edge( VSOURCE , i+1 , 0 , g );
				}
			} else {
				add_edge( i+1 , neighbours[j] , cost , g );
			}
		}
	}

	if ( VERBOSE ) {
		cout << "...Done..." << endl;
	}

	return g;
}

void copy_weight_map( graph_t g , map<pair<int, int>, double> & weights ) {
	weight_map_t weight_map = get( edge_weight, g );
	pair<vertex_iter, vertex_iter> vi;
	for ( vi = boost::vertices( g ) ; vi.first != vi.second ; vi.first++ ) {
		graph_traits<graph_t>::out_edge_iterator out_it, out_it_end;
		tie( out_it, out_it_end ) = out_edges( *vi.first , g );
		for ( ; out_it != out_it_end ; out_it++ ) {
			vertex_t t = target( *out_it , g );
			double w = get(weight_map, *out_it);
			pair<int, int> key(*vi.first,t);
			weights.insert(make_pair(key, w));
		}
	}
}

void copy_distance_map( graph_t g , distance_map_t distance_map , vector_property_map<double> & distances ) {
	pair<vertex_iter, vertex_iter> vi;
	for ( vi = vertices( g ) ; vi.first != vi.second ; vi.first++ ) {
		distances[*vi.first] = distance_map[*vi.first];
	}
}

void copy_weight_and_distance_map( graph_t & g , map<pair<int, int>,double> & weights , distance_map_t & distance_map, vector_property_map<double> & distances ) {
	weight_map_t weight_map = get( edge_weight, g );
	pair<vertex_iter, vertex_iter> vi;
	for ( vi = vertices( g ) ; vi.first != vi.second ; vi.first++ ) {
		distances[*vi.first] = distance_map[*vi.first];
		graph_traits<graph_t>::out_edge_iterator out_it, out_it_end;
		tie( out_it, out_it_end ) = out_edges( *vi.first , g );
		for ( ; out_it != out_it_end ; out_it++ ) {
			vertex_t t = target( *out_it , g );
			double w = get(weight_map, *out_it);
			pair<int, int> key(*vi.first,t);
			weights.insert(make_pair(key, w));
		}
	}	
}

location intToLocation( int loc , int numOfLocations , int pomX ) {
	int ts = (int) floor( ((double) loc-1) / numOfLocations );
	loc = ( ( loc - 1 ) % numOfLocations ) + 1;
	int locx = ((loc-1) % pomX )+1;
	int locy = (( loc - locx ) / pomX )+1;
	
	location l;
	l.ts = ts;
	l.x = locx;
	l.y = locy;
	return l;
}

void drawResults( vector<vector<location>> paths, int scale , int numOfLocations , int numOfFrames , int beginFrame , string gtPath , int pomX , int pomY ) {
	util::Window win("POM_tracker");
	for ( int i = 0 ; i < numOfFrames ; i++ ) {
		img::ImageRGB im ((pomX+1)*scale, (pomY+1)*scale, img::Triple<>());
		for ( unsigned int j = 0 ; j < paths.size() ; j++ ) {
			vector<location> path = paths[j];
			double colorScale = (paths.size()- (double) j) / paths.size();
			for ( unsigned int k = 1 ; k < path.size()-1 ; k++ ) {
				//int loc = path[k];
				//int ts = (int) floor( ((double) loc-1) / numOfLocations );
				location l = path[k];
				if ( l.ts == i ) {
					img::drawEllipse<int>(im, l.x * scale, l.y * scale, 3, 3, 0, img::Triple<>(int(colorScale*255),int(255*(1-colorScale)),int(colorScale*255)), 1);
 					break;
				}
			}
		}
		// draw ground truth
		show_gt( im , beginFrame + i , gtPath , scale );
		cout << "Frame Number: " << i << " ,Frame: " << beginFrame + i << endl;
		win.showImage(im);
		win.waitForKey();
	}
}

void savePaths( vector<vector<location>> paths , int pomX , unsigned int numOfFrames, vector<vector<int>> & fullPaths , string outputPaths ) {
	
	for( unsigned int i = 0 ; i < paths.size() ; i++ ){
		vector<int> path;
		for( unsigned int j = 0 ; j < numOfFrames ; j++ ) {
			path.push_back(0);
		}
		fullPaths.push_back(path);
	}

	for( unsigned int i = 0 ; i < paths.size() ; i++ ) {
		vector<location> path = paths[i];
		for( unsigned int j = 1 ; j < path.size()-1 ; j++ ) {
			int loc = (path[j].y-1) * pomX + path[j].x;
			int ts = path[j].ts;
			//int loc = ( ( path[j] - 1 ) % numOfLocations ) + 1;
			//int ts = (int) floor( ((double) path[j]-1) / numOfLocations );
			fullPaths[i][ts] = loc;
		}
	}

	ofstream outFile;
	string fileLoc = outputPaths + "paths.txt";
	outFile.open(fileLoc.c_str());
	for( unsigned int i = 0 ; i < fullPaths.size() ; i++ ) {
		vector<int> path = fullPaths[i];
		for( unsigned int j = 0 ; j < path.size() ; j++ ) {
			outFile << path[j] << " ";
		}
		outFile << endl;
	}
	outFile.close();
}

void getFilesPerCamera( const string& directoryName, vector< vector< string > >& files ) {
	namespace bf = boost::filesystem;
	bf::path dirPath( directoryName, bf::native );
	if ( !bf::exists( dirPath ) )
		return;
	
	const boost::regex my_filter( "(07|14|19)_([0-9]+)_[0-9]+\\.[0-9]+\\.raw" );
	bf::directory_iterator itrEnd;
	for ( bf::directory_iterator itr( dirPath ); itr != itrEnd; ++itr )
	{
		if ( !bf::is_directory( *itr ) )
		{
			boost::smatch match;
			
			string& fileName = itr->leaf();

			if( boost::regex_match( fileName, match, my_filter ) )
			{
				int cam;
				
				// match camera number
				try
				{
					string s_match(match[1].first, match[1].second);
					cam = lexical_cast<int>(s_match); // string to int conversion					
				}
				catch(const bad_lexical_cast &)
				{
					// ok, couldn't find the camera number, trying to interpret an empty string
					// continue with next file
					continue;
				}

				switch( cam )
				{
					case 14: files[0].push_back( fileName ); break;
					case 19: files[1].push_back( fileName ); break;
					case 7: files[2].push_back( fileName ); break;
				}
			}			
		}
	}
}

void parseLine( string line , vector<vector<rectangle>> & pomLocations ){
	size_t recPos = line.find("RECTANGLE");
	if ( recPos == string::npos ) {
		// contains no rectangle information
		return;
	} else {
		rectangle r;
		// contains rectangle information
		size_t firstSpacePos = line.find( " " );
		size_t secondSpacePos = line.find( " " , firstSpacePos + 1 );
		size_t thirdSpacePos = line.find( " " , secondSpacePos + 1 );
			
		string cam = line.substr( firstSpacePos + 1 , secondSpacePos - firstSpacePos - 1 );
		string loc = line.substr( secondSpacePos + 1 , thirdSpacePos - secondSpacePos - 1 );
		string rec = line.substr( thirdSpacePos + 1 , line.size());

		int camNr = lexical_cast<int>( cam );
		int locNr = lexical_cast<int>( loc );

		if ( rec == "notvisible" ){
			r.isVisible = false;
			r.xmin = -1;
			r.xmax = -1;
			r.ymin = -1;
			r.ymax = -1;
		} else {
			r.isVisible = true;
			size_t xminLoc = rec.find( " " );
			size_t yminLoc = rec.find( " " , xminLoc + 1 );
			size_t xmaxLoc = rec.find( " " , yminLoc + 1 );

			string xminStr = rec.substr( 0 , xminLoc );
			string yminStr = rec.substr( xminLoc + 1 , yminLoc - xminLoc - 1 );
			string xmaxStr = rec.substr( yminLoc + 1 , xmaxLoc - yminLoc - 1 );
			string ymaxStr = rec.substr( xmaxLoc + 1 , rec.size() );

			r.xmin = lexical_cast<int>( xminStr );
			r.ymin = lexical_cast<int>( yminStr );
			r.xmax = lexical_cast<int>( xmaxStr );
			r.ymax = lexical_cast<int>( ymaxStr );
		}

		pomLocations[ camNr ][ locNr ] = r;
	}

}

void readPomFile( string pomFile , vector<vector<rectangle>> & pomLocations ){
	ifstream file(pomFile.c_str());

	if(!file.is_open()){
		cout << "Error opening POM file." << endl;
		exit(1);
	} else {
		string line;
		while(getline(file, line)){
			parseLine( line , pomLocations );
		}
		file.close();
	}
}

img::Triple<> getHeatmapColor( double value ) {
	int r, g, b;

	if ( value > 0.75 && value <= 1 ) {
		r = 255;
		b = 0;
		g = int((1-((value-0.75) * 4)) * 255);
	} else if ( value > 0.5 ) {
		g = 255;
		b = 0;
		r = int(((value-0.5)*4) * 255);
	} else if ( value > 0.25 ) {
		r = 0;
		g = 255;
		b = int((1-((value-0.25) * 4)) * 255);
	} else if ( value > 0 ) {
		r = 0;
		b = 255;
		g = int(value*4*255);
	}

	return img::Triple<>( r , g , b );
}

int calculateSurface( rectangle r ) {
	return ( r.xmax - r.xmin ) * ( r.ymax - r.ymin );
}

img::Triple<> calcAvgColor( img::ImageRGB im ){
	int totalR = 0;
	int totalG = 0;
	int totalB = 0;
	int notCounted = 0;
	for( int i = 0 ; i < im.height() ; i++ ) {
		for ( int j = 0 ; j < im.width() ; j++ ) {
			img::Triple<> color = im(j,i);
			if( color != img::Triple<>( 0, 0, 0 ) && color != img::Triple<>( 255, 255, 255 ) ){
				totalR += color[0];
				totalG += color[1];
				totalB += color[2];
			} else {
				notCounted++;
			}
		}
	}
	int totalPixels = ( im.height() * im.width() ) - notCounted;
	double avgR = (double) totalR / (double) totalPixels;
	double avgG = (double) totalG / (double) totalPixels;
	double avgB = (double) totalB / (double) totalPixels;
	return img::Triple<>( (int) avgR , (int) avgG, (int) avgB );
}

void saveImages( vector<vector<int>> paths, vector<vector<double>> probabilities , string imagePath , string pomFile, int beginFrame , int numOfFrames , int numOfLocations , int pomX, string outputPath , string backgroundPath ){
	cout << "Writing images to disc." << endl;
	int imageWidth = 752;
	int imageHeight = 560;
	const int numOfCameras = 3;
	const string BOG_DIR = "C:/Data/Uva/Scriptie/Code/Cassandra2/config/camera_calib/";

	vector< img::Triple<> > colors;
	colors.push_back(img::Triple<>(255,0,0));	// red
	colors.push_back(img::Triple<>(0,255,0));	// green
	//colors.push_back(img::Triple<>(0,0,255));	// blue
	colors.push_back(img::Triple<>(255,255,0)); // yellow
	colors.push_back(img::Triple<>(0,255,255));	// cyan
	colors.push_back(img::Triple<>(255,0,255)); // purple

	double rot [] = { -20.3973 , -5.82713 , -2.24835 };

	vector< util::Array2D< double > >			distMapX_;
	vector< util::Array2D< double > >			distMapY_;

	distMapX_.resize( numOfCameras, util::Array2D< double >( imageWidth, imageHeight ) );
	distMapY_.resize( numOfCameras, util::Array2D< double >( imageWidth, imageHeight ) );
	
	calib::BogData bogData[numOfCameras];
	cas::help::readBogData( BOG_DIR, bogData );

	for( int cam = 0; cam < numOfCameras ; ++cam )
	{
		cas::help::computeUndistortMap( bogData[cam], distMapX_[cam], distMapY_[cam] );
	}

	// TODO clean up this method, pass on variables and remove hard coded variables;
	vector< vector< string > > filesPerCamera( numOfCameras );
	getFilesPerCamera( imagePath , filesPerCamera );
	vector< img::ImageRGB > inputIms( numOfCameras , img::ImageRGB( imageWidth, imageHeight ) );
	
	vector<vector<rectangle>> pomLocations;
	for (int i = 0 ; i < numOfCameras ; i++ ) {
		vector<rectangle> toAdd (numOfLocations);
		pomLocations.push_back(toAdd);
	}
	readPomFile( pomFile , pomLocations );
	boost::progress_display progress( numOfFrames );
	for ( int i = 0 ; i < numOfFrames ; i++ ) {
		
		vector< img::ImageRGB > imageVec;

		int frame = beginFrame + i;
		string frameNumber;
		if ( frame < 1000 ) {
			frameNumber = "0" + lexical_cast<string>(frame);
		} else {
			frameNumber = lexical_cast<string>(frame);
		}

		for( int cam = 0; cam < numOfCameras; ++cam ) {
			img::ImageG imRaw( imageWidth, imageHeight );
			img::io::readRAW( imagePath + "/" + filesPerCamera[cam][ beginFrame + i ], imRaw, imRaw.width(), imRaw.height() ); 
			img::proc::demosaicPixelGrouping( imRaw, inputIms[cam] );
			
			cas::help::undistortImage( inputIms[cam], distMapX_[cam], distMapY_[cam] );
			img::rotateSafe< img::InterpolatorBilinear >( inputIms[cam], imageWidth/2, imageHeight/2, rot[cam] );

			img::ImageRGB im = inputIms[cam];
			
			img::ImageRGB background;
			string backgroundImPath = backgroundPath + "/result-f" + lexical_cast<string>(frame) + "-c" + lexical_cast<string>(cam) + ".png";
			img::io::read(backgroundImPath , background );

			img::Triple<> maskColor(0,204,153);


			// TODO: ADD THIS TO REVERT BACK!
			
			for ( unsigned int j = 0 ; j < pomLocations[cam].size() ; j++ ) {
				rectangle r = pomLocations[cam][j];
				double value = probabilities[i][j];
				img::Triple<> color = getHeatmapColor( value );
				if ( r.isVisible ) {
					img::drawCross( im , int((r.xmin+r.xmax) / 2.0) , r.ymax, color, 3, 3 );
				}
			}
			
			// AND REMOVE THIS
			/*
			//util::Window win("Cut out");
			for ( unsigned int j = 0 ; j < pomLocations[cam].size() ; j++ ) {
				rectangle r = pomLocations[cam][j];
				if ( r.isVisible ) {
					img::ImageRGB imageCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::copy( im , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , imageCutOut);
					img::ImageRGB backgroundCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::copy( background , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , backgroundCutOut );
					
					img::copyImageForeground( imageCutOut, backgroundCutOut, maskColor, imageCutOut );
					//win.showImage(imageCutOut);
					//win.waitForKey();

					//copyImageForeground( const Image<T, AP1>& imBg, const Image<T, AP2>& imFg, const T& mask_val, Image<T, AP3>& imOut );
					img::Triple<> avgCol = calcAvgColor( imageCutOut );
					img::drawCross( im , int((r.xmin+r.xmax) / 2.0) , r.ymax, avgCol, 3, 3 );
				}
			}
			*/
			//
		
			vector<pair<int,int>> toSort;

			for ( unsigned int j = 0 ; j < paths.size() ; j++ ) {
				int loc = paths[j][i];
				if ( loc != 0 ) {
					rectangle r = pomLocations[cam][loc-1];
					if ( r.xmax != -1 ) {
						int surface = calculateSurface(r);
						toSort.push_back(pair<int,int>(surface,j));
					}
				}
			}
			sort (toSort.begin(), toSort.end());

			vector<cutOutInfo> toAdd;

			for ( unsigned int j = 0 ; j < toSort.size() ; j++ ) {
				int loc = paths[toSort[j].second][i];
				if ( loc != 0 ) {
					rectangle r = pomLocations[cam][loc-1];
					img::ImageRGB backgroundCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::ImageRGB imageCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::copy( background , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , backgroundCutOut );

					img::equals( backgroundCutOut, maskColor, colors[toSort[j].second], img::Triple<>(0,0,0) );

					img::copy( im , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , imageCutOut);
					overlay(imageCutOut,backgroundCutOut, 0.5, true);
					
					cutOutInfo coi;
					coi.r = r;
					coi.pathNr = toSort[j].second;
					coi.imageCutOut = imageCutOut;

					toAdd.push_back(coi);
					
				}
			}

			for ( unsigned int j = 0 ; j < toAdd.size() ; j++ ) {
				cutOutInfo coi = toAdd[j];
				img::copy( coi.imageCutOut, im, coi.r.xmin, coi.r.ymin );
				img::drawBox( im, coi.r.xmin, coi.r.ymin, coi.r.xmax, coi.r.ymax, colors[coi.pathNr]);
			}

			imageVec.push_back( im );
			
		}
		
		++progress;

		img::ImageRGB bigIm( imageVec[0].width() * 2 , imageVec[0].height() * 2 );
		img::copy( imageVec[0] , bigIm , 0 , 0 );
		img::copy( imageVec[1] , bigIm , imageVec[1].width() , 0 );
		img::copy( imageVec[2] , bigIm , imageVec[2].width() / 2 , imageVec[2].height() );

		img::resample( bigIm , imageVec[0].width() , imageVec[0].height() , true );

		string filename = outputPath + frameNumber + ".jpg";
		img::io::write( filename , bigIm );
	}

}

void combinePaths( vector<vector<int>> partialPaths , vector<vector<location>> & completePaths , int numOfLocations , int pomX , int batchSize , int lastTimeStep ) {
	if ( completePaths.size() == 0 ) {
		for ( unsigned int i = 0 ; i < partialPaths.size() ; i++ ) {
			vector<int> partialPath = partialPaths[i];
			vector<location> toAdd;
			for( unsigned int j = 0 ; j < partialPath.size() ; j++ ) {
				location l = intToLocation(partialPath[j] , numOfLocations , pomX );
				toAdd.push_back(l);
			}
			completePaths.push_back(toAdd);
		}
	} else {
		for ( unsigned int i = 0 ; i < partialPaths.size() ; i++ ) {
			int extensionFound = -1;
			vector<int> partialPath = partialPaths[i];
			int l = partialPath[1];
			location loc = intToLocation(l , numOfLocations , pomX );
			vector<location> completePath;
			if ( loc.ts == 0 ) {
				// path starts at beginning of batch
				for ( unsigned int j = 0 ; j < completePaths.size() ; j++ ) {
					completePath = completePaths[j];
					location lastLoc = completePath[completePath.size() - 2];
					if ( lastLoc.ts == lastTimeStep && lastLoc.x == loc.x && lastLoc.y == loc.y ) {
						extensionFound = j;
						break;
					}
				}
			}
			if ( extensionFound != -1 ) {
				completePaths.erase( completePaths.begin() + extensionFound );
				location last = completePath[ completePath.size() - 1 ];
				int ts = last.ts;
				completePath.pop_back();
				for ( unsigned int j = 2 ; j < partialPath.size() ; j++ ) {
					location toAdd = intToLocation( partialPath[j] , numOfLocations , pomX );
					toAdd.ts = ts;
					ts++;
					completePath.push_back(toAdd);
				}
				completePaths.push_back(completePath);
			} else {
				vector<location> toAdd;
				for( unsigned int j = 0 ; j < partialPath.size() ; j++ ) {
					location l = intToLocation(partialPath[j] , numOfLocations , pomX );
					l.ts += lastTimeStep;
					toAdd.push_back(l);
				}
				completePaths.push_back(toAdd);
			}
		}
	}
}

int main(int argc, char* argv[])
{
	// SCENARIO_NAME also defines folder results will be saved in
	string SCENARIO_NAME = "CAS_11-1_500_NEW_IM";
	
	// Path to output folder of POM detector containing .dat files
	//string PROBA_PATH("C:/Data/Uva/Scriptie/Code/pom/result_party08_prior");		// MM
	//string PROBA_PATH("C:/Data/Uva/Scriptie/Code/pom/result_speeddate8-1");		// MM
	string PROBA_PATH("H:/cas_pom/pom_data/scenario11-1/results_200");				// CAS

	// Path to original .raw images
	string IMAGE_PATH("H:/cas_data/scenario11-1");									// CAS

	// Path to pom results .png images
	string BACKGROUND_PATH("H:/cas_pom/pom_data/scenario11-1/results_200");

	// Location of the .pom file, used to later plot the paths onto the video images
	//string POM_FILE("C:/Data/Uva/Scriptie/Code/pom/mm.pom");						// MM
	string POM_FILE("H:/cas_pom/pom_data/scenario11-1/cass.pom");					// CAS

	// size of the POM grid
	//int POM_X = 28;																// MM
	//int POM_Y = 28;																// MM
	int POM_X = 38;																	// CAS
	int POM_Y = 60;																	// CAS

	// Path to GT data, .dat files with locations of the gt objects in pom (non-rounded)coordinats
	string GT_PATH("C:/Data/Uva/Scriptie/Scratch/Matlab/gt-party08");

	int BEGIN_FRAME = atoi(argv[1]);
	int NUM_OF_FRAMES = atoi(argv[2]);
	int BATCH_SIZE = atoi(argv[3]);

	int NUM_OF_BATCHES = int ( ceil( ( double(NUM_OF_FRAMES) / double(BATCH_SIZE) ) ) );

	string OUTPUT_PATH = "C:/Data/Uva/Scriptie/Code/POM_tracker/output/" + SCENARIO_NAME + "/";

	if( ! boost::filesystem::exists( OUTPUT_PATH ) ) {
		boost::filesystem::create_directory( OUTPUT_PATH );
	}

	VERBOSE = false;

	int LOCATIONS_IN_TIMESLICE = POM_X * POM_Y;
	// TOTAL_LOCATIONS is + 2 because of the addition of vsource and vsink
	int TOTAL_LOCATIONS = LOCATIONS_IN_TIMESLICE * NUM_OF_FRAMES + 2;
	
	int DRAW_SCALE = 10;
	
	// Maximum number of people to look for in the scene (i.e. max number of loops the program makes)
	int MAX_DETECTIONS = 25;
	
	// source node has the very first label
	VSOURCE = 0;
	
	// sets if the graph being made is the first one, in that case vsource is 
	// connected to all nodes in the first timeslice
	bool firstGraph = true;

	vector<vector<location>> completePaths;

	int lastTimeStep = -1;

	int batchNr = 0;

	vector<vector<double>> probabilities;

	for( int beginFrame = BEGIN_FRAME ; beginFrame < BEGIN_FRAME + NUM_OF_FRAMES ; beginFrame += BATCH_SIZE ) {
		batchNr++;
		cout << "Batch nr. " << batchNr << " of " << NUM_OF_BATCHES 
			<< "." << endl;
		int numOfFrames;
		if ( beginFrame + BATCH_SIZE > BEGIN_FRAME + NUM_OF_FRAMES ) {
			numOfFrames = BEGIN_FRAME + NUM_OF_FRAMES - beginFrame;
		} else {
			numOfFrames = BATCH_SIZE;
		}

		int startFrame = beginFrame;
		if ( beginFrame != BEGIN_FRAME ) {
			startFrame--;
			numOfFrames++;
		}

		vector<location> sources;
		if ( completePaths.size() != 0 ) {
			for( unsigned int i = 0 ; i < completePaths.size() ; i++ ) {
				vector<location> path = completePaths[i];
				location l = path[path.size() - 2];
				if( l.ts == lastTimeStep ) {
					sources.push_back(l);
				}
			}
		}
		
		int source;
		if ( sources.size() == 0 ) {
			source = VSOURCE;
		} else {
			location l = sources[0];
			source = (l.y-1) * POM_X +l.x;
			sources.erase(sources.begin());
		}
		

		// Variables
		distance_map_t distance_map; 
		vector<vector<int>> paths, prevPaths;

		// Create graph
		vector<double> vertices = readFiles( startFrame , numOfFrames, PROBA_PATH , probabilities );
		int vsink = vertices.size() + 1;
		graph_t g_original;
		createGraph( vertices , vsink, LOCATIONS_IN_TIMESLICE , g_original , POM_X , firstGraph , completePaths , BATCH_SIZE , lastTimeStep );
		firstGraph = false;
		// Find shortest path using dag shortest paths method
		graph_t g;
		g = g_original;
		vector<int> path = find_shortest_path_dag( g , source, vsink , distance_map );
		
		if( VERBOSE ) {
			printPath(path , LOCATIONS_IN_TIMESLICE , POM_X );
		}
		paths.push_back(path);

		cout << paths.size() << " path found." << endl;

		double costCurrent = paths_cost( paths, g_original );
		double costPrev = 0.0;
		
		vector_property_map<double> distances;
		map<pair<int, int>, double> weights;

		copy_weight_and_distance_map( g , weights , distance_map , distances );

		for( int i = 1 ; i <= MAX_DETECTIONS ; i++ ) {
			if( i != 1 ) {
				if ( sources.size() == 0 ) {
					costCurrent = paths_cost( paths, g_original );

					if ( VERBOSE ) {
						cout << "Current cost is:   " << costCurrent << endl;
						cout << "Previous cost was: " << costPrev << endl;
					}
					
					if( costCurrent >= costPrev ) {
							break;
					}
				}
				cout << paths.size() << " paths found." << endl; 
			}
			
			costPrev = costCurrent;
			prevPaths = paths;
			
			graph_t g = g_original;
			map<int,int> vertexPairs;
			extend_graph( g , paths , vsink, vertexPairs );

			transform_edge_cost( g , weights , distances , vsink , TOTAL_LOCATIONS, vertexPairs );

			if ( sources.size() == 0 ) {
				source = VSOURCE;
			} else {
				location l = sources[0];
				source = (l.y-1) * POM_X +l.x;
				sources.erase(sources.begin());
			}
			
			// Find shortest path using Dijkstras shortest path method
			path = find_shortest_path_dijkstra( g , source , vsink , distance_map );

			if( VERBOSE ) {
				printPath(path, LOCATIONS_IN_TIMESLICE , POM_X );
			}

			interlace_and_augment( path , paths , g_original );

			weights.clear();
			
			copy_weight_and_distance_map( g , weights , distance_map , distances );
		}

		if ( VERBOSE ) {
			//printPaths( prevPaths , LOCATIONS_IN_TIMESLICE , POM_X );
		}

		if ( batchNr != NUM_OF_BATCHES ) {
			probabilities.pop_back();
		}

		combinePaths( prevPaths , completePaths , LOCATIONS_IN_TIMESLICE , POM_X  , BATCH_SIZE , lastTimeStep );
		lastTimeStep += BATCH_SIZE;

	}
	vector<vector<int>> convertedPaths;
	savePaths( completePaths , POM_X , NUM_OF_FRAMES, convertedPaths , OUTPUT_PATH );

	saveImages( convertedPaths , probabilities , IMAGE_PATH , POM_FILE , BEGIN_FRAME , NUM_OF_FRAMES , LOCATIONS_IN_TIMESLICE, POM_X , OUTPUT_PATH , BACKGROUND_PATH);

	//drawResults( completePaths , DRAW_SCALE , LOCATIONS_IN_TIMESLICE , NUM_OF_FRAMES , BEGIN_FRAME , GT_PATH , POM_X , POM_Y );

}